

import tensorflow as tf


def simple_loop():
    sess = tf.Session()
    i = tf.constant(0)

    kk = tf.constant(12)
    # c = lambda i: tf.less(i, 10)
    # b = lambda i: tf.add(i, 1)
    # r = tf.while_loop(c, b, [i])
    # print(sess.run(r))

    # def cond(i,kk):
    #     return i<10
    c = lambda i,kk : i<10
    def body(i,kk):
        kk = tf.add(kk,5)
        i = tf.Print(i + 1, ["I:", i ," KK: ", kk ])
        return ( i , kk )
    ret = tf.while_loop(c, body, [i,kk])
    print(sess.run(ret))


def concat_loop():
    sess = tf.Session()
    i0 = tf.constant(0)
    m0 = tf.ones([2, 2])
    c = lambda i, m: i < 10
    b = lambda i, m: [i + 1, tf.concat([m, m], axis=0)]
    i_out , m_out = tf.while_loop(
        c, b, loop_vars=[i0, m0],
        shape_invariants=[i0.get_shape(), tf.TensorShape([None, 2])])
    i_ , m_ = sess.run([i_out,m_out])
    print("i:{}, m:{}".format(i_,m_.shape))

if __name__ == '__main__':
    # simple_loop()
    concat_loop()